# -*- coding: utf-8 -*-
'''
@Time    : 2022/8/10 14:04
@Author  : HandsomeWu
@File    : dataSplit.py
@Describe: 数据集的切分
        训练集包含n个样本，每个样本有一个样本编号和一个类别编号。
        假设某一类别的样本共有m个，则将编号最小的m/2(向上取整)个样本作为训练集，将其他样本作为测试集
* input:
    两个正整数n和k，分别表示样本数量和类别数量
    n个正整数，第i个正整数j表示第i个样本的类别编号是j

* output
    输出包含两行
    第一行是训练集样本编号，编号从小到大输出 中间用空格隔开
    第二行是测试集的样本编号，编号从小到大输出，中间用空格隔开。
'''
import math

def data_split(n,k,data):
    # Step 1 :先计算出最少类别的 m/2是多少

    data_count = {}
    for i in data:
        if i not in data_count:
            data_count[i] = 1
        else:
            data_count[i] +=1

    split_point = math.ceil(min(data_count.values()) / 2)
    # Step 2: 分割数据集
    split_set = dict(zip(data_count.keys(),[split_point]*k))
    train_list = []
    test_list = []

    for idx,i in enumerate(data):
        if split_set[i] == 0:
            test_list.append(idx+1)
        else:
            split_set[i] -= 1
            train_list.append(idx+1)

    return train_list,test_list







if __name__ == '__main__':
    N,K =10,3
    dataset = [3,2,2,1,2,3,1,3,3,3]

    train,test  = data_split(N,K,dataset)
    print(train,test)